import os
import requests
import zipfile
from io import BytesIO
from PIL import Image
from time import sleep
from random import uniform

REGIONS = ["africa", "asia", "europe"]
MALE_COUNT = 100
FEMALE_COUNT = 200
OUTPUT_DIR = "avatars"

def download_image(category, idx):
    try:
        response = requests.get("https://thispersondoesnotexist.com/image", timeout=10)
        if response.status_code == 200:
            img = Image.open(BytesIO(response.content)).resize((256, 256))
            folder = os.path.join(OUTPUT_DIR, category)
            os.makedirs(folder, exist_ok=True)
            path = os.path.join(folder, f"{category}_{idx:03d}.jpg")
            img.save(path)
            print(f"[+] Saved {path}")
            return True
    except Exception as e:
        print(f"[-] Failed {category} #{idx}: {e}")
    return False

def download_category(gender, region, count):
    cat = f"{gender}_{region}"
    for i in range(count):
        success = False
        retry = 0
        while not success and retry < 3:
            success = download_image(cat, i + 1)
            if not success:
                retry += 1
                sleep(uniform(1.0, 2.0))
            else:
                sleep(uniform(0.5, 1.2))

def zip_folder(folder, zip_name):
    with zipfile.ZipFile(zip_name, "w") as zipf:
        for root, _, files in os.walk(folder):
            for file in files:
                full_path = os.path.join(root, file)
                arcname = os.path.relpath(full_path, folder)
                zipf.write(full_path, os.path.join("avatars", arcname))

def main():
    print("[*] Starting download...")
    for region in REGIONS:
        download_category("male", region, MALE_COUNT)
        download_category("female", region, FEMALE_COUNT)

    print("[*] Zipping avatars...")
    zip_folder(OUTPUT_DIR, "world_avatars.zip")
    print("[✔] Done! Output: world_avatars.zip")

if __name__ == "__main__":
    main()
